#from tensorflow import keras as K
from tensorflow.contrib import keras as K
from data_handle.read_json import chat_history_read
import model.intent_classify.cnn_net_ori as cnn_net
from config.model_parameter import model_parameter
import time
import os
import tensorflow as tf

def train_cnn_net(model, x_pad, y_pad,
                  x_pad_eval, y_pad_eval,
                  num_epochs, batch_size, model_prefix):

    cp_callback = K.callbacks.ModelCheckpoint(model_prefix + '.ckpt',
                                              save_weights_only=True,
                                              verbose=1)
    model.fit(x_pad,
              y_pad,
              epochs=num_epochs,
              batch_size=batch_size,
              validation_data=(x_pad_eval, y_pad_eval),
              verbose=1,
              callbacks=[cp_callback])

    model_prefix = "data/model_path/" + model_prefix
    model.save_weights(model_prefix + '_weights.hdf5')
    model.save(model_prefix + '.h5')
    return model


def load_cnn_net(model_path):
    model = K.models.load_model(model_path)
    return model


def cnn_net_predict(model, id_to_cat, x_pad):
    y_proba = model.predict(x_pad)
    y_id = y_proba.argmax()
    predict_label = [ id_to_cat[each] for each in y_id]
    return predict_label



if __name__ == '__main__':

    model_parameter_ins = model_parameter()
    experiment_name = model_parameter_ins.flags.FLAGS.experiment_name
    FLAGS = model_parameter_ins.get_parameter(experiment_name).FLAGS

    #model = load_cnn_net(model_path = "data/model_path/test.h5")

    #管理GPU
    start_time = time.time()
    # Config GPU options
    if FLAGS.per_process_gpu_memory_fraction == 0.0:
        gpu_options = tf.GPUOptions(allow_growth=True)
    elif FLAGS.per_process_gpu_memory_fraction == 1.0:
        gpu_options = tf.GPUOptions()

    else:
        gpu_options = tf.GPUOptions(
            per_process_gpu_memory_fraction=FLAGS.per_process_gpu_memory_fraction)

    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.cuda_visible_devices
    chat_history_read_ins = chat_history_read(FLAGS)
    chat_history_read_ins.simple_concat()
    # chat_history_read_ins.statistics()
    X_train, X_test, y_train_intent_label, y_test_intent_label, \
    y_train_sub_task, y_test_sub_task, y_train_detail, y_test_detail\
        = chat_history_read_ins.get_train_data()
    # X_train = [X_train,chat_history_read_ins.sub_mask,chat_history_read_ins.detail_mask]
    # X_test  = [X_test,chat_history_read_ins.sub_mask,chat_history_read_ins.detail_mask]

    y_train = [y_train_intent_label,y_train_sub_task,y_train_detail]
    y_test =  [y_test_intent_label,y_test_sub_task,y_test_detail]
    char_size = chat_history_read_ins.embeddings_matrix.shape[0]
    embedding_dim = chat_history_read_ins.embeddings_matrix.shape[1]

    # model = cnn_net.create_net(embedding_dim = embedding_dim,
    #                            char_size = char_size,
    #                            num_filters = 5,
    #                            embeddings_matrix = chat_history_read_ins.embeddings_matrix,
    #                            hidden_dim = 100,
    #                            input_length = FLAGS.input_length,
    #                            sub_mask=chat_history_read_ins.sub_mask,
    #                            detail_mask=chat_history_read_ins.detail_mask)

    model = cnn_net.create_net(embedding_dim = embedding_dim,
                               char_size = char_size,
                               num_filters = 5,
                               embeddings_matrix = chat_history_read_ins.embeddings_matrix,
                               hidden_dim = 100,
                               input_length = FLAGS.input_length)


    train_cnn_net(model = model,
                  x_pad = X_train,
                  y_pad = y_train,
                  x_pad_eval = X_test,
                  y_pad_eval = y_test,
                  num_epochs = 1,
                  batch_size = 32,
                  model_prefix = "nlp_class_ori")

